import matplotlib
import matplotlib.pyplot as plt

paNode = dict(boxstyle="sawtooth",fc="0.8")
leafNode = dict(boxstyle="round4",fc="0.8")
arrow_args = dict(arrowstyle = "<-")

def plotNode(nodeTxt,centerPt,parentPt,nodeType):
    createPlot.axl.annotate(nodeTxt,xy = parentPt,xycoords='axes fraction',xytext = centerPt,ha="center",bbox = nodeType,arrowprops=arrow_args)

def createPlot():
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    createPlot.axl = plt.subplot(111,frameon = False)
    plotNode('paNode',(0.5,0.1),(0.1,0.5),paNode)
    plotNode('leafNode',(0.8,0.1),(0.3,0.8),leafNode)
    plt.show()

def getNumLeafs(myTree):
    '''
    获取叶子节点，确定x轴的情况
    '''
    numLeafs = 0
    firstStr = list(myTree.keys())[0]
    second = myTree[firstStr]
    for k in second.keys():
        if type(second[k]).__name__ == 'dict':
            numLeafs += getNumLeafs(second[k])
        else:
            numLeafs += 1
    return numLeafs

def getTreeDepth(myTree):
    '''
    获取树的深度,确定y轴的情况
    '''
    depth = 0
    firstStr = list(myTree.keys())[0]
    second = myTree[firstStr]
    for key in second.keys():
        if type(second[k]).__name__ == 'dict':
            thisdepth = 1 + getTreeDepth(second[k])
        else:
            thisdepth = 1
        if thisdepth > depth: depth = thisdepth
    return depth

def retrieveTree(i):
    listOfTrees = [{'no surfacing': {0: 'no',1: {'flippers': {0: 'no',1: 'yes'}}}},
    {'no surfacing': {0: 'no',1: {'flippers': {0: {'head': {0: 'no',1: 'yes'}},1: 'no'}}}}]
    return listOfTrees[i]
if __name__ == '__main__':
    mytree = retrieveTree(0)
    a = getNumLeafs(mytree)
    print(a)